use std::collections::{hash_map::Entry, HashMap};

/// 并查集算法
struct UnionFind {
    ancestors: HashMap<&'static str, &'static str>,
}
impl UnionFind {
    pub fn new() -> Self {
        Self {
            ancestors: HashMap::new(),
        }
    }

    /// 查找祖先节点，同时进行路径折叠
    pub fn find_ancestor(&mut self, x: &'static str) -> &'static str {
        match self.ancestors.entry(x) {
            Entry::Occupied(entry) => {
                let y = *entry.get();
                if x == y {
                    x
                } else {
                    let z = self.find_ancestor(y);
                    self.ancestors.insert(x, z);
                    z
                }
            }
            Entry::Vacant(entry) => {
                entry.insert(x);
                x
            }
        }
    }

    /// 将 x 的祖先节点设置为 y，需要确保 y 的祖先是自己
    pub fn merge_to(&mut self, x: &'static str, y: &'static str) {
        let xf = self.find_ancestor(x);
        if xf != y {
            self.ancestors.insert(xf, y);
        }
    }

    /// 统计集合数量
    pub fn count_sets(&self) -> usize {
        self.ancestors.iter().filter(|(k, v)| k == v).count()
    }
}

pub fn count_provinces() -> String {
    let mut testcase = 0;
    let mut uf = UnionFind::new();
    let mut answers = String::new();

    for line in include_str!("../district.json").lines() {
        if line.starts_with("    \"") || line.starts_with("}") {
            testcase += 1;
            if testcase > 1 {
                answers.push_str(&uf.count_sets().to_string());
                answers.push(',');
            }
            uf = UnionFind::new();
            continue;
        }

        let (main, others) = match line.split_once(':') {
            Some(data) => data,
            None => continue,
        };
        let main = main.trim().trim_matches('"');
        let others = others
            .trim()
            .trim_matches(|c| "[],".contains(c))
            .split(',')
            .map(|it| it.trim().trim_matches('"'));

        let parent = uf.find_ancestor(main);
        for other in others {
            uf.merge_to(other, parent);
        }
    }

    answers.pop();
    answers
}
